#include "refine/smooth_despill_c.h"
#include "smooth_despill_impl.h"
#include "math/pi.h"

namespace refine {
namespace smooth_despill {
namespace c {

//everything, which is below or equal to alpha0 is considered as 0,
//and so for alpha1: x >= alpha1 => x = uchar_max
const uint alpha0 = uchar_max / 20;
const uint alpha1 = uchar_max - alpha0;

/// contains multitude of unit vectors, which lie in one half-space
struct unitvecs
{
	std::vector<vectorf> v;
	unitvecs(uint N = 40) 
	{
		v.reserve(N * N);
		for (uint a = 0; a < N; ++a) {
			for (uint b = 0; b < N; ++b) {
				float alpha = PI * float(a) / N;		//[0, pi)
				float beta = PI * float(b) / N - PI_2;	//[-pi/2, pi/2)
				float cb = cos(beta); // cb >= 0 always
				v.push_back(vectorf(
					cb * sin(alpha), // always >= 0
					cb * cos(alpha),
					sin(beta)
				));
			}
		}
	}
};

/// contains color statistics for either background or foreground
class summator
{
public:
	summator()
		{ clear(); }
	void clear() {
		N = 0;
		I = vectoru::nill;
		II.zero();
	}
	summator & operator += (vectoru const & i) {
		++N;
		I += i;
		II += tsquare(i);
		return * this;
	}
	summator & operator -= (vectoru const & i) {
		--N;
		I -= i;
		II -= tsquare(i);
		return * this;
	}
	uint count() const
		{ return N; }
	vectorf mean() const
		{ return (1.0f / N) * vectorf(I); }
	smatrixf covariance() const
		{ return (1.0f / (N - 1)) * (smatrixf(II) - (1.0f / N) * tsquare(vectorf(I))); }
private:
	uint N; ///< the count of pixels
	vectoru I; ///< sum of I
	smatrixu II; ///< sum of I*I^T
};

/// stores statistics for foreground and background in two different summators
class summator_pair
{
public:
	summator_pair()
		{ clear(); }
	void clear() {
		s0.clear();
		s1.clear();
	}
	template <typename Color>
	summator_pair & operator += (Color const & c) {
		if (c.a <= alpha0)
			s0 += v(c);
		else if (c.a >= alpha1)
			s1 += v(c);
		return * this;
	}
	template <typename Color>
	summator_pair & operator -= (Color const & c) {
		if (c.a <= alpha0)
			s0 -= v(c);
		else if (c.a >= alpha1)
			s1 -= v(c);
		return * this;
	}
	/// calculates means and inversed common covariance matrices for both foreground and background,
	/// returns false if condition number is too bad for that calculation
	bool distributions(float max_condition_number, vectorf & mu0, vectorf & mu1, smatrixf & iCov0, smatrixf & iCov1, float * condition_number = 0) const
	{
		if (s0.count() < 2 || s1.count() < 2 ||
			s0.count() * max_condition_number < s1.count() ||
			s1.count() * max_condition_number < s0.count())
			return false;

		if (condition_number)
			*condition_number = (s1.count() >= s0.count()) ? 
				float(s1.count()) / s0.count() : float(s0.count()) / s1.count();

		smatrixf S0 = s0.covariance();
		smatrixf S1 = s1.covariance();
		float max0, max1;
		S0.eigenvalues(0, 0, &max0);
		S1.eigenvalues(0, 0, &max1);
		float max = std::max(max0, max1);
		if (max <= 0)
			return false;
		S0 += 0.001f * max;
		S1 += 0.001f * max;

		mu0 = s0.mean();
		mu1 = s1.mean();
		iCov0 = S0.inverse();
		iCov1 = S1.inverse();

		return true;
	}
private:
	summator s0, s1;
};

template <typename Raster>
class alpha_smoother : public alpha_smoother_base
{
public:
	alpha_smoother(
		Raster const & im_, 
		Raster & next_im_, ///< it receives updates in alpha channel
		img::bitmap & modifiable, 
		std::vector<vectorf> const & unitvecs_,
		float max_condition_number/*, float more_contrast*/)
		: alpha_smoother_base(modifiable, max_condition_number/*, more_contrast*/), 
		  im(im_), next_im(next_im_), unitvecs(unitvecs_)
	{
		clear(); 
	}
	void clear() {
		remove_all();
		checks = achanges = 0;
	}
	void add(size_point const & p)
		{ s += im[p]; }
	void remove(size_point const & p)
		{ s -= im[p]; }
	void remove_all() 
		{ s.clear(); }
	void process(size_point const & p, size_rect const &) 
	{
		++checks;

		vectorf mu0, mu1;
		smatrixf A0, A1;
		float condition_number;
		if (!s.distributions(max_condition_number, mu0, mu1, A0, A1, &condition_number)) {
			modifiable[p] = false;
			return;
		}

		vectorf I = v(im[p]);
		vectorf d0 = I - mu0;
		vectorf d1 = I - mu1;
		vectorf Ad0 = A0(d0);
		vectorf Ad1 = A1(d1);

		vectorf n;
		float E = 0;
		for (std::vector<vectorf>::const_iterator i = unitvecs.begin(); 
			i != unitvecs.end(); ++i) 
		{
			float e = 
				square(scalar_product(*i, Ad0)) / scalar_product(*i, A0(*i)) +
				square(scalar_product(*i, Ad1)) / scalar_product(*i, A1(*i));
			if (e > E) {
				E = e;
				n = *i;
			}
		}
		assert (E > 0);

		float t0 = - scalar_product(n, Ad0) / scalar_product(n, A0(n));
		float t1 = - scalar_product(n, Ad1) / scalar_product(n, A1(n));
		unsigned alpha = classify_end(condition_number, t0, t1);

		if (alpha == im[p].a) {
			modifiable[p] = false;
			return;
		}

		next_im[p].a = alpha;
		++achanges;
	}
private:
	Raster const & im;
	Raster & next_im;
	std::vector<vectorf> const & unitvecs;
	summator_pair s;
};

template <typename Raster>
class despiller : public despiller_base<Raster>
{
public:
	despiller(
		Raster const & im_,	Raster & next_im_, 
		float max_condition_number)
		: despiller_base<Raster>(im_, next_im_, max_condition_number)
	{
		clear(); 
	}
	void clear() {
		remove_all();
	}
	void add(size_point const & p)
		{ s += im[p]; }
	void remove(size_point const & p)
		{ s -= im[p]; }
	void remove_all() 
		{ s.clear(); }
	void process(size_point const & p, size_rect const &) 
	{
		vectorf mu0, mu1;
		smatrixf A0, A1;
		if (!s.distributions(max_condition_number, mu0, mu1, A0, A1)) {
			//next_im[p].a = (im[p].a <= uchar_avg) ? 0 : uchar_max;
			return;
		}

		vectorf I = v(im[p]);
		vectorf d0 = I - mu0;
		vectorf d1 = I - mu1;
		vectorf Ad0 = A0(d0);
		vectorf Ad1 = A1(d1);

		uint a = im[p].a; //255 * alpha
		uint b = uchar_max - a; //255 * (1 - alpha)

		//solve the system M*n=r
		smatrixf M = float(a * a) * A0 + float(b * b) * A1;
		vectorf r = float(a) * Ad0 - float(b) * Ad1;
		vectorf n = M.inverse() * r; //n is less in uchar_max times...
		vectorf c1 = I + float(b) * n; //... that's why no additional division is required here
		discrete(c1, next_im[p]);
	}
private:
	summator_pair s;
};

template <typename Raster>
static void perform_core(
	Raster & im, 
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Smooth-X & Despill" << endl;

	size_extent dim = im.dim();
	assert (dim == modifiable.dim());
	unitvecs uvecs;

	Raster despilled = im;

	//smooth alpha passes
	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		alpha_smoother<Raster> sm(im, despilled, modifiable, uvecs.v, p.max_condition_number/*, p.more_contrast*/);
		window_pass(sm, p.whs);

		if (!p.silent)
			clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;

		if (sm.change_count() > 0)
			copy_alpha(im, despilled);
		else 
			break;

		if (pass+1 < p.passes_num)
			img::dilate8(modifiable, p.whs);
	}

	// no smoothing, despill only
	despiller<Raster> de(im, despilled, p.max_condition_number);
	window_pass(de, p.whs);
	im = despilled;
}

template <typename Raster>
static void perform_from_grey_mask(
	Raster & im, 
	const params & p)
{
	size_extent dim = im.dim();

	img::bitmap not_min(dim), not_max(dim);
	for (typename Raster::const_iterator i = im.begin(); i != im.end(); ++i) {
		not_min[i] = (*i).a != 0;
		not_max[i] = (*i).a != uchar_max;
	}

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(im, modifiable, p);
}

// exported functions

void perform(
	byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, p);
}

void perform(
	argb_byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, p);
}


} //namespace c
} //namespace smooth_despill
} //namespace refine
